{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/sdb1/hy17/venv/lib/python3.6/site-packages/numba/errors.py:131: UserWarning: Insufficiently recent colorama version found. Numba requires colorama >= 0.3.9\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\" # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "# Importing the libraries \n",
    "import torch\n",
    "import torchaudio\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import matplotlib.cm as cm\n",
    "import pandas as pd\n",
    "import os\n",
    "import librosa\n",
    "import librosa.display\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import time\n",
    "# from Models import Autoencoder\n",
    "# from TestData import SingleMus\n",
    "from Data_TIMIT import Data_TIMIT\n",
    "import IPython.display as ipd\n",
    "from IPython.display import clear_output\n",
    "from Blocks import BasicBlock, Bottleneck, ChannelChange\n",
    "# from EntropyControl_curr import AE_control\n",
    "from Prop_sr import Prop_Model\n",
    "from utils import *\n",
    "from validation import *\n",
    "\n",
    "# Checking the availability of CUDA on machine\n",
    "# print(\"Torch\", torch.__version__, \"CUDA\", torch.version.cuda)\n",
    "# print(\"Device:\", torch.device(\"cuda:3\"))\n",
    "# print(torch.cuda.is_available())\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- #### TIMIT Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = 'example_data/train_5db.pth'\n",
    "test_path = 'example_data/test_5db.pth'\n",
    "train_loader = torch.load(train_path)\n",
    "test_loader = torch.load(test_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3000 300\n"
     ]
    }
   ],
   "source": [
    "print(len(train_loader),len(test_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Define"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "if train_path == 'data/half_win_train_ctn_std4_m.pth':\n",
    "    db = '_5db'\n",
    "else:\n",
    "    db = '_0db'\n",
    "filters = 30\n",
    "d = 6\n",
    "f2 = 60\n",
    "m = 160\n",
    "sr = True\n",
    "lr = 0.0001\n",
    "weight1 = 1/5\n",
    "weight2 = 1/60\n",
    "target = 5\n",
    "ratio = 1/2\n",
    "label = time.strftime(\"%m%d_%H%M%S\")\n",
    "model_name = label+'ratio_sr40_d'+str(d)+db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(30, 6, 60, 160, True(5,0.50), 0.2, 0.016) - 0604_023532ratio_sr40_d6_0db \n"
     ]
    }
   ],
   "source": [
    "hypara = '({}, {}, {}, {}, {}({},{:.2f}), {}, {}) - {} '\\\n",
    "      .format(filters, d, f2, m, sr, target, ratio, str(weight1)[:5], str(weight2)[:5],\\\n",
    "             model_name)\n",
    "print(hypara)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hypara = '({}, {}, {}, {}, {}({},{}), {}) - {} '\\\n",
    "#       .format(filters, d, f2, m, sr, target_clean, target_noise, weight,\\\n",
    "#              model_name)\n",
    "# print(hypara)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bottleneck\n"
     ]
    }
   ],
   "source": [
    "model = Prop_Model(block = Bottleneck, scale = 10, filters = filters, d_s = d, d_n = d,\\\n",
    "                   f2 = f2, num_m = m, sr = sr, ratio = ratio).cuda()\n",
    "\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(list(model.parameters()) + [model.means_s] + [model.means_n] \\\n",
    "                             + [model.ratio], lr=lr)\n",
    "t = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 50\n",
    "epoch_list = []\n",
    "max_sdr = (0,0,0,0)\n",
    "max_rs = None\n",
    "flct = 0.1\n",
    "etp = 0\n",
    "itermax = 0\n",
    "# model.stage = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while 1: \n",
    "    start = time.time()\n",
    "    model.train()\n",
    "    if t == 3:\n",
    "        model.stage = 1\n",
    "    if model.scale < 500:\n",
    "        model.scale *= t+1\n",
    "    train_sdr_s = []\n",
    "    train_sdr_rs = []\n",
    "    train_sdr_rx = []\n",
    "    \n",
    "    train_etp_cl = []\n",
    "    train_etp_ns = []\n",
    "    control = 0\n",
    "    \n",
    "    for wave_s, wave_x, source, mixture in train_loader:\n",
    "        \n",
    "        source = source[0]\n",
    "        mixture = mixture[0]\n",
    "        \n",
    "        noise = (mixture - source).cuda()\n",
    "        source = source.cuda()\n",
    "        mixture = mixture.cuda()\n",
    "\n",
    "        s_h, n_h, prob_s, prob_n = model(mixture, soft = True)\n",
    "        \n",
    "        entp_cl = None\n",
    "        entp_ns = None\n",
    "        if not isinstance(prob_s, type(None)) and not isinstance(prob_n, type(None)):\n",
    "            entp_cl = entropy_prob(prob_s)\n",
    "            train_etp_cl.append(entp_cl.cpu().data.numpy())\n",
    "            entp_ns = entropy_prob(prob_n)\n",
    "            train_etp_ns.append(entp_ns.cpu().data.numpy())    \n",
    "            \n",
    "\n",
    "        loss = criterion(s_h, source) + criterion(n_h+s_h, mixture)\n",
    "        \n",
    "        if etp == 1:\n",
    "            control = 1\n",
    "            loss += weight1 * ((target - entp_cl - entp_ns)**2) + \\\n",
    "                    weight2 * ((ratio - entp_ns/entp_cl)**2)\n",
    "       \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    end = time.time()\n",
    "    print('epoch_{}| Time:{:.0f} | Control:{:.0f} | Stage: {} | Itermax: {} | etp:{} '\\\n",
    "          .format(t, end-start, control, model.stage, itermax, etp))\n",
    "\n",
    "#     s_s_score, s_rs_score, s_rx_score, s_diff_score, _, _, _ = test_newData(True)\n",
    "    h_s_score, h_rs_score, h_rx_score, h_diff_score, max_single_rs,  arg_s, arg_n = test(False, test_loader, model)\n",
    "    \n",
    "    entropy = [0, 0]\n",
    "\n",
    "    if not isinstance(arg_s, type(None)):\n",
    "        entropy[0] = cal_entropy(arg_s.data.cpu().numpy().flatten())\n",
    "    if not isinstance(arg_n, type(None)):\n",
    "        entropy[1] = cal_entropy(arg_n.data.cpu().numpy().flatten())\n",
    "      \n",
    "    epoch_list.append((h_s_score, h_rs_score, h_diff_score, entropy[0]))\n",
    "    print('|Test-hard s: {:.2f}, rs: {:.2f}, rx: {:.2f}, diff: {:.2f} \\n\\\n",
    "          |Entropy : {:.2f}, {:.2f}'\\\n",
    "          .format(h_s_score, h_rs_score, h_rx_score, h_diff_score, \n",
    "                  entropy[0], entropy[1]))\n",
    " #          |Test-soft s: {:.2f}, rs: {:.2f}, rx: {:.2f}, diff: {:.2f} \\n\\\n",
    "    \n",
    "    t += 1\n",
    "    if model.stage >= 1:\n",
    "        itermax += 1\n",
    "        if h_rx_score > max_sdr[1]:\n",
    "            max_sdr = (h_rs_score, h_rx_score, entropy[0], entropy[1])\n",
    "            max_rs = max_single_rs\n",
    "            torch.save(model, model_name+'.model')  \n",
    "            itermax = 0\n",
    "    \n",
    "    if model.stage == 1 and etp == 0 and itermax >= 3:\n",
    "        etp = 1\n",
    "        itermax = 0\n",
    "        max_sdr = (0,0,0,0)\n",
    "\n",
    "# #     if model.stage == 1 and itermax >= 2:\n",
    "#     if model.stage == 1 and etp == 1 and itermax >= 3:\n",
    "#         model.stage = 2  \n",
    "#         itermax = 0\n",
    "# #         max_sdr = (0,0,0,0)\n",
    "#         print('Enter stage 2')\n",
    "\n",
    "    if etp == 1 and itermax > 10:\n",
    "        print('over')\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_name = model_name+str(max_sdr[1])[:5]\n",
    "output = '({}, {}, {}, {}, {}({},{:.2f}), {}, {}) - {}-({:.2f},{:.2f})({:.2f},{:.2f}) '\\\n",
    "      .format(filters, d, f2, m, sr, target, ratio, str(weight1)[:5], str(weight2)[:5],\\\n",
    "             model_name, max_sdr[0], max_sdr[1], max_sdr[2], max_sdr[3])\n",
    "print(output)\n",
    "with open('output_result.txt', 'a') as the_file:\n",
    "    the_file.write(output+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = max_rs[1].cpu().data.numpy()[0]\n",
    "rx = max_rs[1].cpu().data.numpy()[0]+max_rs[2].cpu().data.numpy()[0]\n",
    "librosa.output.write_wav(model_name+'_rs.wav', rs/max(rs), sr = 16000)\n",
    "librosa.output.write_wav(model_name+'_rx.wav', rx/max(rx), sr = 16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(max_rs[0].cpu().data.numpy(), rate = 16000)\n",
    "# librosa.output.write_wav(model, save_max_rs_1[0].cpu().data.numpy()[0], sr = 16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(rs/max(rs), rate = 16000)\n",
    "# librosa.output.write_wav('0508_whole_sp32_1392rs.wav', rs/max(rs), sr = 16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(rx/max(rx), rate = 16000)\n",
    "# librosa.output.write_wav('0508_whole_sp32_1392rx.wav', rx/max(rx), sr = 16000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
